%% This code runs Stage I and II of the method.
close all; clear; 
warning('off')
clc

%% define saved files from running this code.
estFile_save = 'est1_output_2006.mat';
fig1_save = 'fig1_CatchSSB_fit.png';
fig2_save = 'fig2_F_heatmap.png';
figC1_save = 'figC1_age_polynomial.png';
figC2_save = 'figC2_fishing_mortality.png';
figC3_save = 'figC3_N0_distribution.png';

%% load data
load 'ATE_data_age10.mat'

% set time period for years up to Tend
Tend = 2015; % 2006 (main result) or 2015 (full sample)
ATE_catch = ATE_catch(year<=Tend);
SSB = SSB(year<=Tend);
Recruit = Recruit(year<=Tend);
F_at = F_at(year<=Tend,:);
year = year(year<=Tend);

ttmax = length(year); % total number of years
% F_at, Recruit and SSB start in t0 = 1968
t0 = year(find(Recruit>0,1)); 

H = (1e+3)*ATE_catch; % harvest (kg)
SSB = (1e+3)*SSB; % biomass (kg)

% normalize the data series
normH = norm(H); normSSB = norm(SSB(year >= t0)); normF = norm(F_at(year>=t0,:));
H0 = H./normH; SSB0 = SSB./normSSB;
F_at0 = F_at./normF;

%% set model parameters
run set_parameters.m

%% Stage I -- biomass parameters
% (a) initialize (R,lambda,f)
rng('default') % for replication

R0_scale = rand()/rand(); lambda0 = rand();
f0 = zeros(ttmax,1); % fishing mortality
est0 = [R0_scale; lambda0; f0];
myoption= optimset('Display','off','MaxFunEvals',1e+10,...
    'TolX',1e-10,'MaxIter',1e+8);

fprintf('(a) initialize (R,lambda,f) by fitting catch and SSB:\n')
est1 = fminunc('catchSSB_dist',est0,myoption, ...
    H0,SSB0,normH,normSSB,Recruit,T,w_a,m_a,mature);
fprintf('... completed.\n')

% (b) OLS estimate Theta for h(a,K=4)
K = 4; % degree of polynomial
f = est1(3:end)*ones(1,10); % f(t) estimated from (a)
fprintf('(b) initialize polynomial h(a) with degree K=%d\n',K)
y = F_at(year>=t0,:)./f(year>=t0,:);
a = ones(size(y,1),1)*(1:10);
y = y(:); a = a(:);
X_age = [a a.^2 a.^3 a.^4 a.^5 a.^6 a.^7];
x = X_age(:,1:K);
ols_fit = fitlm(x,y);
Theta0 = ols_fit.Coefficients.Estimate;
fprintf('... completed.\n')
clear f ols_fit

% (c) estimate best-fitting (R,lambda,ft,Theta)
fprintf('(c) solve the best-fitting parameters ... completed.\n')
% initialization
est2 = [est1;Theta0];
[est,fval,e,o,g,h] = fminunc('catchSSBF_dist',est2,myoption, ...
        H0,SSB0,normH,normSSB,F_at0,normF,Recruit,T,w_a,m_a,mature);

% point estimate ----
R1_scale = est(1); lambda = est(2);
ft = est(3:(ttmax+2));
Theta = est((ttmax+3):end);

% standard errors are the sqrt of the diagonal of the inv Hessian
std_err = sqrt(diag(inv(h)));
se_Rscale = std_err(1);
se_lambda = std_err(2);
se_ft = std_err(3:(ttmax+2));
se_Theta = std_err((ttmax+3):end);

%% Stage II -- catchability calibration
% use number of vessels as a proxy for effort
boat = [508; 910]; % number of vessels in 2014 and 2015
fend2 = ft(end-1:end); % estimated fishing mortality in t and t-1
q0 = mean(fend2./boat);
E = (1/q0)*ft; % generate effort: E(t)

%% save estimates from model run
save(estFile_save,'E','R1_scale','lambda','ft','q0',...
    'se_Rscale','se_lambda','se_ft','se_Theta','Theta','K')

%% Plot: polynomial h(a) for age = 1 to 10
XX_age = [ones(T,1) age' (age.^2)' (age.^3)' (age.^4)' (age.^5)' ...
    (age.^6)' (age.^7)'];
X_age = XX_age(:,1:(K+1));
h_age = X_age(:,1:K+1)*Theta;

fig_ha = figure('Position',[0 0 700 400]);
plot(age,h_age,'LineWidth',3)
xlabel 'age, a'
ylabel 'value of function, h(a)'
xlim([1 10])

%% Plot: fishing mortality f(t) estimates
fig_ft = figure('Position',[0 0 700 600]);
subplot(2,1,1)
plot(year,ft,'LineWidth',1.5)
title '(a) Estimated fishing mortality time tend, f_t'
xlabel 'Year, t'
ylabel 'value of f_t'
xlim([1950 2015])

subplot(2,1,2)
plot(year,se_ft,':','LineWidth',2)
title '(b) Standard Error of f_t'
xlabel 'Year, t'
ylabel 'standard error'
xlim([1950 2015])

%% Plot: distribution of initial abundance
R1 = max(Recruit)*R1_scale;
N1 = R1*exp(-lambda*age');
fig_N0 = figure('Position',[0 0 700 400]);
bar(0:T,[R1;N1])
title ({'Estimated distribution for abundance in year 0';
    ['Recruitment = ' num2str(R1_scale) '*R_{max} (se = ' num2str(se_Rscale) ')'];
    ['\lambda = ' num2str(lambda) ' (se = ' num2str(se_lambda) ')']})
xlabel 'Age, a'
ylabel 'Number of Fish, N_{a,0}'
xlim([0 T])


%% compute predicted catch and SSB
% extrapolate F (age 1-50)
h_age(11:end) = h_age(10);
F_hat = ft*h_age';
z = F_hat + ones(ttmax,1)*m_a(1:T)'; % total mortality (t,a)

% population matrix simulation ---------%
N = zeros(ttmax,T+1); % row= year, col=age
N(year < t0,1) = R1;
N(year >=t0,1) = Recruit(year>=t0);
N(1,2:end) = N1;
for t = 2:ttmax
    for j = 2:(T+1)
        N(t,j) = N(t-1,j-1)*exp(-z(t-1,j-1));
    end % for j
end % for t
N_end = N(end,:); % end-year population age structure

% predicted catch & SSB ----------------------%
% age-structured catch H(t,a)
H_age = (F_hat./z).*(1-exp(-z)).*(ones(ttmax,1)*w_a).*N(:,2:end);
H_hat = H_age*ones(T,1);

% predicted spawning stock biomass
SSB_hat = N(:,2:end)*(mature.*(w_a'));

%% Plot: fitting of Catch and SSB 
fig_catchSSB = figure('Position',[0 0 700 570]);
subplot(2,1,1)
hold on
plot(year,(1e-3)*H,'Color',[0.9294    0.6902    0.1294],...
    'LineWidth',4,'DisplayName','record')
plot(year,(1e-3)*H_hat,'k-.','LineWidth',2.4,'DisplayName','predict')
title '(a) Predicted Catch vs. Recorded Catch'
lgd = legend('show','Location','NE');
lgd.EdgeColor = 'none';
xlabel 'Year'
ylabel 'Catch (metric tons)'
xlim([1950 2015])
ylim([0 inf])
hold off

subplot(2,1,2)
hold on
SSB2 = SSB; SSB2(SSB <=0) = nan;
plot(year,(1e-3)*SSB2,'Color',[0.9294    0.6902    0.1294],...
    'LineWidth',4,'DisplayName','record')
plot(year,(1e-3)*SSB_hat,'k-.','LineWidth',2.4,'DisplayName','predict')
title '(b) Predicted SSB vs. Recorded SSB'
lgd = legend('show','Location','NE');
lgd.EdgeColor = 'none';
xlabel 'Year'
ylabel 'SSB (metric tons)'
xlim([1950 2015])
ylim([0 inf])
hold off

%% save H_hat and SSB_hat (for appendix Fig E.1)
% save('est1_predict_HSSB_2015.mat','H','H_hat','SSB2','SSB_hat')
%% Plot: heatmap of F
fig_F = figure('Position',[0 0 900 1000]);
subplot(1,2,1)
heatmap(1:10,1968:Tend,F_at(year>=t0,:));
c = gray;c = flipud(c);
colormap(c)
brighten(0.5)
title 'age-structured fishing mortality from ICCAT'
xlabel 'age 1-10'
subplot(1,2,2)
heatmap(1:10,1968:Tend,F_hat(year>=t0,1:10));
c = gray;c = flipud(c);
colormap(c)
brighten(0.5)
title(sprintf('predicted fishing mortality with h(a,K=%d)',K))
xlabel 'age 1-10'

%% save output figures
saveas(fig_catchSSB,fig1_save)
saveas(fig_F,fig2_save)
saveas(fig_ha,figC1_save)
saveas(fig_ft,figC2_save)
saveas(fig_N0,figC3_save)